#include <pthread.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <stdarg.h>
#include <netpacket/packet.h>
#include <sys/socket.h>
#include <sys/ioctl.h>
#include <arpa/inet.h>
#include <net/if.h>
#include <netinet/in.h>
#include <netinet/tcp.h>
#include <net/ethernet.h>
#include <time.h>

#include <unistd.h>
#include <fcntl.h>

#include "common.h"

#define DL_RATE_PROTO 0x8811
#define BUF_SIZE 2000
#define PAYLOAD_SIZE 1000-14-20
#define USAGE "\
sender -if [interface_name_for_sending] -dip [destination_ip_addr]\
 -stp [source_transport_port] -dtp [destination_transport_port]\n"


void* rate_ctrl_msg(void* spec);
int get_ip_for_interface(char* intf_name);
int str_matches( const char* given, int num_args, ... );
void verbose_println( const char* format, ... );
long checksum(unsigned short *addr, unsigned int count);
void create_tcp_packet(char* buffer, uint16_t len, uint32_t s_ip, uint32_t d_ip, unsigned short* tcp);

int verbose;

int main(int argc, char** argv) {
	pthread_t tid;
	struct in_addr iaddr;
	uint32_t dst_ip;
	uint16_t src_tp;
	uint16_t dst_tp;
	struct sender_spec_t spec;	
	verbose =0;
	int i;
	unsigned short tcp_buf[2000];

	for (i=1; i<argc || argc<=1; i++) {
		if (argc<=1 || str_matches(argv[i], 5, "-?", "-help", "--help", "help", "?")) {
			printf(USAGE);
			return 0;
		}
		if (str_matches(argv[i], 1, "-v")) {
			verbose = 1;
		}
                if (str_matches(argv[i], 1, "-if")) {
                        i = i + 1;
                        if (i==argc) {
                                printf("-if needs a name...die\n");
                                return -1;
                        }
                        spec.src_ip = (uint32_t)get_ip_for_interface(argv[i]);
                }
		if (str_matches(argv[i], 1, "-dip")) {
			i = i + 1;
			if (i==argc) {
				printf("-dip needs an ip address...die\n");
				return -1;
			}
			if (inet_aton(argv[i], &iaddr) == 0) {
				printf("invalid ip address...die\n");
				return -1;
			}
			spec.dst_ip = (uint32_t)iaddr.s_addr;
		}
                if (str_matches(argv[i], 1, "-stp")) {
                        i = i + 1;
                        if (i==argc) {
                                printf("-stp needs a port number...die\n");
                                return -1;
                        }
                        spec.src_tp = (uint16_t)atoi(argv[i]);
                }
                if (str_matches(argv[i], 1, "-dtp")) {
                        i = i + 1;
                        if (i==argc) {
                                printf("-dtp needs a port number...die\n"); 
                                return -1;
                        }
                        spec.dst_tp = (uint16_t)atoi(argv[i]);
                }
	}
	spec.rate = 35000;

	/* 
	 * OPEN SOCKET FOR SENDING 
	 */

	int snd_fd = socket(AF_INET, SOCK_RAW, IPPROTO_TCP);
	if (snd_fd < 0) {
		printf("can not create socket\n");
		return -1;
	}

	int val = 0;
	if(setsockopt(snd_fd, IPPROTO_IP, IP_HDRINCL, &val, sizeof(val)) < 0) {
		printf("setsocket IP_HDRINCL error\n");
		return -1;
	}

	struct sockaddr_in addr;
	addr.sin_family = AF_INET;
	addr.sin_port = 0;
	addr.sin_addr.s_addr = spec.src_ip;
	if(bind(snd_fd, (struct sockaddr*)&addr, sizeof(addr)) != 0) {
                printf("bind error\n");
                return -1;
	}

	addr.sin_addr.s_addr = spec.dst_ip;
	if(connect(snd_fd, (struct sockaddr*)&addr, sizeof(addr)) != 0) {
		printf("can not connect to ip address %u", spec.dst_ip);
		return -1;
	}

	printf("Source IP: %u, Destination IP:%u\n", spec.src_ip, spec.dst_ip);

	/* 
	 * spawn a thread for receiving rate controller msges
	 */

	if( pthread_create(&tid, NULL, rate_ctrl_msg, &spec) != 0 ) {
                printf("can't spawn rate_ctrl_msg thread...die\n");
		return -1;
        }
	
	
	/*
	 * Prepare TCP header
	 */
	char* send_buffer[PAYLOAD_SIZE];

	struct tcphdr* mytcp = (struct tcphdr*)send_buffer;
	mytcp->source = htons(spec.src_tp);
	mytcp->dest = htons(spec.dst_tp);
	mytcp->doff = 5;
	mytcp->window = htons(16000);

	uint64_t sleep_time;	// in nano second
	struct timespec rqtp;
	struct timespec rmtp;
	/*
	 *  the sending job done here 
	 */
	uint32_t seq = 0;
	verbose_println("Start the sending process");
	while(1) {
		mytcp->seq = seq;
		create_tcp_packet(send_buffer, PAYLOAD_SIZE, spec.src_ip, spec.dst_ip, tcp_buf);
		
                uint64_t sec = 8000 / spec.rate;
                uint64_t nsec = (uint64_t)(8000 * 1e9 / spec.rate - sec*1e9);
		//verbose_println("rate is %lld. sleeping for %lld sec and %lld nsec", 
		//		spec.rate, sec, nsec);
		rqtp.tv_sec = sec;
		rqtp.tv_nsec = nsec;
		nanosleep( &rqtp, &rmtp);
		write(snd_fd, send_buffer, PAYLOAD_SIZE);
		seq += 1;
		//verbose_println("Sender sent a packet of 1000B. current rate: %u",spec.rate);
	}
	return 0;
}

void* rate_ctrl_msg(void* input) {
        char* buffer[BUF_SIZE];
	sender_spec_t* spec = (sender_spec_t*) input;
        RateMsg* castedRM;

	printf("rate control thread starts to run. rate cntrl msg type=%u\n",DL_RATE_PROTO);
        int fd = socket(AF_PACKET, SOCK_RAW, htons(DL_RATE_PROTO));
        if (fd < 0) {    
                printf("Can not create socket fd= %d... Die\n",fd);
                goto err;
        }
        while (1) {
                int rd = read(fd, buffer, BUF_SIZE);
                verbose_println("rate controller read %d bytes\n", rd);
                castedRM = (RateMsg *)buffer; 

		verbose_println("**rate cntrl msg => src_ip: %u ... dst_ip: %u ... src_tp: %u ... dst_tp: %u",
		castedRM->src_ip, castedRM->dst_ip, ntohs(castedRM->src_tp), ntohs(castedRM->dst_tp));
		verbose_println("this flow spec => src_ip: %u ... dst_ip: %u ... src_tp: %u ... dst_tp: %u ***",
		spec->src_ip, spec->dst_ip, spec->src_tp, spec->dst_tp);


		if (castedRM->src_ip == spec->src_ip && castedRM->dst_ip == spec->dst_ip
			&& ntohs(castedRM->src_tp) == spec->src_tp && ntohs(castedRM->dst_tp) == spec->dst_tp) {
			spec->rate = castedRM->rate;
			verbose_println("Set the rate to: %d\n", spec->rate);
		}
	}

err:	return NULL;
}

/*
 * Note: buffer has tcphdr at first
 */
void create_tcp_packet(char* buffer, uint16_t len, uint32_t s_ip, uint32_t d_ip, unsigned short* tcp) {

	struct tcphdr * mytcp = (struct tcphdr*)buffer;
        int tcpdatalen = len - (mytcp->doff*4);

	struct tcp_pseudo pseudohead;
        pseudohead.src_addr=s_ip;
        pseudohead.dst_addr=d_ip;
        pseudohead.zero=0;
        pseudohead.proto=IPPROTO_TCP;
        pseudohead.length=htons(sizeof(struct tcphdr) + tcpdatalen);

        int totaltcp_len = sizeof(struct tcp_pseudo) + sizeof(struct tcphdr) + tcpdatalen;


        memcpy((unsigned char *)tcp,&pseudohead,sizeof(struct tcp_pseudo));
        memcpy((unsigned char *)tcp+sizeof(struct tcp_pseudo),(unsigned char *)mytcp,sizeof(struct tcphdr));
        memcpy((unsigned char *)tcp+sizeof(struct tcp_pseudo)+sizeof(struct tcphdr), 
		(unsigned char *)mytcp+(mytcp->doff*4), tcpdatalen);

        uint16_t chk = (uint16_t)checksum(tcp,totaltcp_len);
	mytcp->check = chk;

}

long checksum(unsigned short *addr, unsigned int count) {
          /* Compute Internet Checksum for "count" bytes
            *         beginning at location "addr".
            */
       register long sum = 0;


        while( count > 1 )  {
           /*  This is the inner loop */
               sum += * addr++;
               count -= 2;
       }
           /*  Add left-over byte, if any */
       if( count > 0 )
               sum += * (unsigned char *) addr;

           /*  Fold 32-bit sum to 16 bits */
       while (sum>>16)
           sum = (sum & 0xffff) + (sum >> 16);

       return ~sum;
}


#define MAX_INTERFACES 32

/** returns the NBO IP address associated with intf_name, or 0 on failure */
int get_ip_for_interface(char* intf_name) {
    struct ifconf ifc;
    struct ifreq ifrp[MAX_INTERFACES];
    int fd, i, ret;

    /* give ifc a buffer to populate */
    ifc.ifc_len = sizeof(ifrp);
    ifc.ifc_buf = (caddr_t)&ifrp;

    /* create a socket so we can use ioctl */
    if((fd = socket(PF_INET, SOCK_DGRAM, 0)) < 0) {
        printf("error in finding ip address for interface\n");
        return 0;
    }

    /* populate ifc with a list of the system's interfaces */
    if(ioctl(fd, SIOCGIFCONF, &ifc) < 0) {
        printf("error in finding ip address for interface (ioctl)\n");
        return 0;
    }

    /* loop over each returned interface info */
    for(i=0; i<ifc.ifc_len / sizeof(struct ifreq); i++) {
        if(ifrp[i].ifr_addr.sa_family == AF_INET) {
            /* is this the interface we were asked about? */
            if( strncmp(ifrp[i].ifr_name, intf_name, IFNAMSIZ)==0 ) {
                ret = ((struct sockaddr_in *)&ifrp[i].ifr_addr)->sin_addr.s_addr;
                break;
            }
        }
        else {
            printf("error: unexpected family for %s (bailing out)\n", ifrp[i].ifr_name);
            ret = 0;
            break;
        }
    }

    /* close our socket and return the address we found */
    close(fd);
    return ret;
}


int str_matches( const char* given, int num_args, ... ) {
    int ret;
    const char* str;
    va_list args;
    va_start( args, num_args );

    ret = 0;
    while( num_args-- > 0 ) {
        str = va_arg(args, const char*);
        if( strcmp(str,given) == 0 ) {
            ret = 1;
            break; /* found a match */
        }
    }
    va_end( args );

    return ret;
}

void verbose_println( const char* format, ... ) {
    va_list args;
    if(!verbose)
        return;

    va_start( args, format );

    vfprintf( stderr, format, args );
    fprintf( stderr, "\n" );

    va_end( args );
}
